from ljp.engine import Trainer


def get_data(dataname, batch_size=64, image_size=64, path_to_your_data=''):
    data = None
    if dataname == 'cifar100':
        from ljp.dataset.databuilder_cifar import CIFAR100
        data = CIFAR100(batch_size=batch_size, dataroot=f'{path_to_your_data}/cifar100/', image_size=image_size)
    if dataname == 'cifar10':
        from ljp.dataset.databuilder_cifar import CIFAR10
        data = CIFAR10(batch_size=batch_size, dataroot=f'{path_to_your_data}/cifar10/', image_size=image_size)
    if dataname == 'svhn':
        from ljp.dataset.databuilder_cifar import SVHN
        data = SVHN(batch_size=batch_size, image_size=image_size)

    if dataname in ['pathmnist', 'dermamnist', 'octmnist', 'bloodmnist', 'tissuemnist', 'organamnist', ]:
        from ljp.dataset.databuilder_medmnist import MEDMNIST
        data = MEDMNIST(flag=dataname, image_size=image_size, batch_size=batch_size)
    return data


def get_model(modelname, data):
    model = None
    if modelname in ['resnet50', 'resnet18', 'resnet34',
                     ]:
        from ljp.models.resnet import resnet18, resnet34, resnet50

        model = {
            'resnet50': resnet50,
            'resnet18': resnet18,
            'resnet34': resnet34,

        }[modelname](num_classes=data.num_classes, inchannel=data.inchannels)
        model = Trainer(model)

    return model
